// --*- C++ -*------x-----------------------------------------------------------
//
//
// Author:          Enrico Negri, Silvio Tosatto
//
// Project name:    BiocompUP Align
//
// Date:            01/2008
//
// Description:     This program calculates sub-optimal alignments for two
//                  sequences. Alignments can be either global, local or
//                  free-shift; both with or without secondary structure.
//
// -----------------x-----------------------------------------------------------

#include <AGPFunction.h>
#include <Alignment.h>
#include <AlignmentBase.h>
#include <AlignmentData.h>
#include <FSAlign.h>
#include <GetArg.h>
#include <NWAlign.h>
#include <ReverseScore.h>
#include <ScoringP2P.h>
#include <ScoringP2S.h>
#include <ScoringS2S.h>
#include <Sec.h>
#include <SecSequenceData.h>
#include <SequenceData.h>
#include <SubMatrix.h>
#include <SWAlign.h>
#include <iostream>

#include <HenikoffProfile.h>
#include <PSICProfile.h>
#include <SeqDivergenceProfile.h>

#include <LogAverage.h>
#include <CrossProduct.h>
#include <DotPFreq.h>
#include <DotPOdds.h>
#include <Pearson.h>
#include <JensenShannon.h>
#include <EDistance.h>


using namespace Biopool;


/// Show command line options and help text.
void
sShowHelp()
{
	cout << "\nSUBOPTIMAL ALIGNMENT GENERATOR"
	     << "\nThis program calculates sub-optimal alignments for two sequences. Alignments can be either global,"
	     << "\nlocal or free-shift; both with or without secondary structure. It is possible to select between"
	     << "\nsequence-to-sequence, profile-to-sequence or profile-to-profile alignments. Profiles (optional)"
	     << "\nare extracted from BLAST M6 output files, which can be further subjected to a number of sequence"
	     << "\nsimilarity filters to optimize similarity.\n"
	     << "\nOptions:"
	     << "\n * [--in <name>]     \t Name of input FASTA file"
	     << "\n   [--pro1 <name>]   \t Name of target profile (BLAST M6 format) file"
	     << "\n   [--pro2 <name>]   \t Name of template profile (BLAST M6 format) file"
	     << "\n   [--sec <name>]    \t Name of secondary structure FASTA file"
	     << "\n   [--out <name>]    \t Name of output FASTA file (default = to screen)"
	     << "\n   [-m <name>]       \t Name of substitution matrix file (default = blosum62.dat)"
	     << "\n   [-M <name>]       \t Name of secondary structure substitution matrix file (default = secid.dat)"
	     << "\n   [--cSeq <double>] \t Coefficient for sequence alignment (default = 1.00)"
	     << "\n   [--cSec <double>] \t Coefficient for secondary structure alignment (default = 1.00)"
	     << "\n   [-p <double>]     \t Penalty multiplier for suboptimal alignments (default = 1.00)"
	     << "\n   [-a <double>]     \t Penalty subtractor for suboptimal alignments (default = 1.00)"
	     << "\n   [-d <double>]     \t Min. master vs all seq. id. filter threshold (default = 0.00)"
	     << "\n   [-D <double>]     \t Min. all vs all seq. id. filter threshold (default = 0.00)"
	     << "\n   [-u <double>]     \t Max. master vs all seq. id. filter threshold (default = 1.00)"
	     << "\n   [-U <double>]     \t Max. all vs all seq. id. filter threshold (default = 1.00)"
	     << "\n   [-o <double>]     \t Open gap penalty (default = 11.00)"
	     << "\n   [-e <double>]     \t Extension gap penalty (default = 3.00)"
	     << "\n   [-n <int>]        \t Number of suboptimal alignments (default = 5)"
	     << "\n   [-r <int>]        \t Reverse sequence Z-score using N data points (default = 0, i.e. off)"
	     << "\n   [--fasta]         \t Use FASTA format to load profiles"
	     << "\n   [--global]        \t Needleman-Wunsch global alignment"
	     << "\n   [--local]         \t Smith-Waterman local alignment"
	     << "\n   [--freeshift]     \t Free-shift alignment (default)"
	     << "\n   [--verbose]       \t Verbose mode"
	     << "\n" << endl;
}


int
main(int argc, char **argv)
{
	// Default parameters

	string matrixFileName = "blosum62.dat";
	string matrixSecFileName = "secid.dat";
	string seq1, seq2, sec1, sec2, seq1Name, seq2Name;
	string inputFileName, proFile1Name, proFile2Name, secFileName, outputFileName;
	double cSeq = 1.00;
	double cSec = 1.00;
	double suboptPenalty, suboptPenaltyAdd;
	double downs, downa, ups, upa;
	double openGapPenalty = 11.00;
	double extensionGapPenalty = 3.00;
	unsigned int suboptNum;
	int nReverse = 0;


	// --------------------------------------------------
	// 0. Treat options
	// --------------------------------------------------

	if (getArg("h", argc, argv))
	{
		sShowHelp();
		return 1;
	}

	getArg("-in", inputFileName, argc, argv, "!");
	getArg("-pro1", proFile1Name, argc, argv, "!");
	getArg("-pro2", proFile2Name, argc, argv, "!");
	getArg("-sec", secFileName, argc, argv, "!");
	getArg("-out", outputFileName, argc, argv, "!");
	getArg("m", matrixFileName, argc, argv, matrixFileName);
	getArg("M", matrixSecFileName, argc, argv, matrixSecFileName);
	getArg("-cSeq", cSeq, argc, argv, cSeq);
	getArg("-cSec", cSec, argc, argv, cSec);
	getArg("p", suboptPenalty, argc, argv, 1.00);
	getArg("a", suboptPenaltyAdd, argc, argv, 1.00);
	getArg("d", downs, argc, argv, 999.9);
	getArg("D", downa, argc, argv, 999.9);
	getArg("u", ups, argc, argv, 999.9);
	getArg("U", upa, argc, argv, 999.9);
	getArg("o", openGapPenalty, argc, argv, openGapPenalty);
	getArg("e", extensionGapPenalty, argc, argv, extensionGapPenalty);
	getArg("n", suboptNum, argc, argv, 5);
	getArg("r", nReverse, argc, argv, nReverse);
	bool fasta = getArg("-fasta", argc, argv);
	bool global = getArg("-global", argc, argv);
	bool local = getArg("-local", argc, argv);
	bool freeshift = getArg("-freeshift", argc, argv);
	bool verbose = getArg("-verbose", argc, argv);

	if (!local && !global)
		freeshift = true;


	// --------------------------------------------------
	// 1. Load data
	// --------------------------------------------------

	string path = getenv("VICTOR_ROOT");
	if (path.length() < 3)
		cout << "Warning: environment variable VICTOR_ROOT is not set." << endl;
	string examplesPath; // = path + "examples/";
	string dataPath = path + "data/";

	if (inputFileName != "!")
	{
		inputFileName = examplesPath + inputFileName;
		ifstream inputFile(inputFileName.c_str());
		if (!inputFile)
			ERROR("Error opening input FASTA file.", exception);
		Alignment ali;
		ali.loadFasta(inputFile);
		if (ali.size() < 1)
			ERROR("Input FASTA file must contain two sequences.", exception);
		seq1 = Alignment::getPureSequence(ali.getTarget());
		seq2 = Alignment::getPureSequence(ali.getTemplate(0));
		seq1Name = ali.getTargetName();
		seq2Name = ali.getTemplateName(0);
	}
	else
		ERROR("subali needs input FASTA file.", exception);

	if (secFileName != "!")
	{
		secFileName = examplesPath + secFileName;
		ifstream secFile(secFileName.c_str());
		if (!secFile)
			ERROR("Error opening secondary structure FASTA file.", exception);
		Alignment aliS;
		aliS.loadFasta(secFile);
		if (aliS.size() < 1)
			ERROR("Secondary structure FASTA file must contain two sequences.", exception);
		sec1 = Alignment::getPureSequence(aliS.getTarget());
		sec2 = Alignment::getPureSequence(aliS.getTemplate());
	}

	matrixFileName = dataPath + matrixFileName;
	ifstream matrixFile(matrixFileName.c_str());
	if (!matrixFile)
		ERROR("Error opening substitution matrix file.", exception);

	matrixSecFileName = dataPath + matrixSecFileName;
	ifstream matrixSecFile(matrixSecFileName.c_str());
	if (!matrixSecFile)
		ERROR("Error opening secondary structure substitution matrix file.", exception);


	// --------------------------------------------------
	// 2. Output test data
	// --------------------------------------------------

	if (verbose)
	{
		fillLine(cout);
		if (secFileName != "!")
			cout << "\nTarget sequence:\n"                                       << seq1
			     << "\n\nTarget secondary structure:\n"                          << sec1
			     << "\n\nTemplate sequence:\n"                                   << seq2
			     << "\n\nTemplate secondary structure:\n"                        << sec2
			     << "\n\nRead substitution matrix from file: "                   << matrixFileName
			     << "\nRead secondary structure substitution matrix from file: " << matrixSecFileName
			     << "\nCoefficient for sequence alignment            = "         << cSeq
			     << "\nCoefficient for secondary structure alignment = "         << cSec
			     << "\nOpen gap penalty      = "                                 << openGapPenalty
			     << "\nExtension gap penalty = "                                 << extensionGapPenalty
			     << "\n" << endl;
		else
			cout << "\nTarget sequence:\n"                     << seq1
			     << "\n\nTemplate sequence:\n"                 << seq2
			     << "\n\nRead substitution matrix from file: " << matrixFileName
			     << "\nOpen gap penalty      = "               << openGapPenalty
			     << "\nExtension gap penalty = "               << extensionGapPenalty
			     << "\n" << endl;
	}
	fillLine(cout);


	// --------------------------------------------------
	// 3. Select alignment mode
	// --------------------------------------------------

	SubMatrix sub(matrixFile);
	SubMatrix subStr(matrixSecFile);
	AGPFunction gf(openGapPenalty, extensionGapPenalty);

	AlignmentData *ad;
	Structure *str;
	ScoringFunction *sf;
	ScoringScheme *ss;

	if (proFile1Name != "!")
	{
		// Construct first profile

		proFile1Name = examplesPath + proFile1Name;
		ifstream proFile1(proFile1Name.c_str());
		if (!proFile1)
			ERROR("Error opening target profile (BLAST M6 format) file.", exception);

		Alignment ali1;
		if (fasta)
			ali1.loadFasta(proFile1);
		else
		{
			cout << "We have for MSA1:" << endl;
			ali1.loadBlastMode6FullSeq(proFile1, seq1);
		}

		if (downs < 999.9)
			ali1.RemoveLowerSimple(downs);
		else
			if (downa < 999.9)
				ali1.RemoveLowerAll(downa);

		if (ups < 999.9)
			ali1.RemoveUpperSimple(ups);
		else
			if (upa < 999.9)
				ali1.RemoveUpperAll(upa);

		Profile *pro1 = new Profile();
		pro1->setProfile(ali1);

		if (proFile2Name != "!")
		{
			// --------------------------------------------------
			// 3.1. Profile-to-Profile case
			// --------------------------------------------------

			// Construct second profile

			proFile2Name = examplesPath + proFile2Name;
			ifstream proFile2(proFile2Name.c_str());
			if (!proFile2)
				ERROR("Error opening template profile (BLAST M6 format) file.", exception);

			Alignment ali2;
			if (fasta)
				ali2.loadFasta(proFile2);
			else
			{
				cout << "We have for MSA2:" << endl;
				ali2.loadBlastMode6FullSeq(proFile2, seq2);
			}

			if (downs < 999.9)
				ali2.RemoveLowerSimple(downs);
			else
				if (downa < 999.9)
					ali2.RemoveLowerAll(downa);

			if (ups < 999.9)
				ali2.RemoveUpperSimple(ups);
			else
				if (upa < 999.9)
					ali2.RemoveUpperAll(upa);

			Profile *pro2 = new Profile();
			pro2->setProfile(ali2);


			// Print MSA

			cout << "FASTA 2" << "\n"
			     << seq2      << "\n"
			     << "MSA 2"   << "\n"
			     << ali2.getTarget() << endl;

			for (unsigned int i = 0; i < ali2.size(); i++)
				cout << ali2.getTemplate(i) << endl;
			cout << endl;


			if (secFileName != "!")
			{
				ad = new SecSequenceData(4, seq1, seq2, sec1, sec2, seq1Name,
					seq2Name);
				str = new Sec(&subStr, ad, cSec);
				sf = new DotPFreq(pro1, pro2);
				ss = new ScoringP2P(&sub, ad, str, pro1, pro2, sf, cSeq);
			}
			else
			{
				ad = new SequenceData(2, seq1, seq2, seq1Name, seq2Name);
				sf = new DotPFreq(pro1, pro2);
				ss = new ScoringP2P(&sub, ad, 0, pro1, pro2, sf, 1.00);
			}
		}
		else
		{
			// --------------------------------------------------
			// 3.2. Profile-to-Sequence case
			// --------------------------------------------------

			if (secFileName != "!")
			{
				ad = new SecSequenceData(4, seq1, seq2, sec1, sec2, seq1Name,
					seq2Name);
				str = new Sec(&subStr, ad, cSec);
				ss = new ScoringP2S(&sub, ad, str, pro1, cSeq);
			}
			else
			{
				ad = new SequenceData(2, seq1, seq2, seq1Name, seq2Name);
				ss = new ScoringP2S(&sub, ad, 0, pro1, 1.00);
			}
		}
	}
	else
	{
		// --------------------------------------------------
		// 3.3. Sequence-to-Sequence case
		// --------------------------------------------------

		if (secFileName != "!")
		{
			ad = new SecSequenceData(4, seq1, seq2, sec1, sec2, seq1Name,
				seq2Name);
			str = new Sec(&subStr, ad, cSec);
			ss = new ScoringS2S(&sub, ad, str, cSeq);
		}
		else
		{
			ad = new SequenceData(2, seq1, seq2, seq1Name, seq2Name);
			ss = new ScoringS2S(&sub, ad, 0, 1.00);
		}
	}


	// --------------------------------------------------
	// 4. Calculate alignments
	// --------------------------------------------------

	Align *a;

	if (global)
	{
		cout << "\nSuboptimal Needleman-Wunsch alignments:\n" << endl;
		a = new NWAlign(ad, &gf, ss);
	}
	else
		if (local)
		{
			cout << "\nSuboptimal Smith-Waterman alignments:\n" << endl;
			a = new SWAlign(ad, &gf, ss);
		}
		else
		{
			cout << "\nSuboptimal free-shift alignments:\n" << endl;
			a = new FSAlign(ad, &gf, ss);
		}


	// --------------------------------------------------
	// 5. Output alignments
	// --------------------------------------------------

	a->setPenalties(suboptPenalty, suboptPenaltyAdd);

	vector<Alignment> a2 = a->generateMultiMatch(suboptNum);
	if (a2.size() == 0)
		ERROR("No output alignments generated.", exception);

	a2[0].cutTemplate(1);
	Alignment a3 = a2[0];

	for (unsigned int i = 1; i < a2.size(); i++)
	{
		a2[i].cutTemplate(1);
		a3.addAlignment(a2[i]);
	}

	if (outputFileName != "!")
	{
		outputFileName = examplesPath + outputFileName;
		ofstream outputFile(outputFileName.c_str());
		if (!outputFile)
			ERROR("Error opening output FASTA file.", exception);
		cout << "Saving output to FASTA file: " << outputFileName << endl;
		a3.saveFasta(outputFile);
	}
	else
		a3.saveFasta(cout);
	cout << endl;
	fillLine(cout);

	return 0;
}
